import numpy as np
from typing import List, Tuple

class NoiseEstimator:
    def __init__(self, fft_size: int, sr: int = 16000, bands_config: str = "webRTC"):
        """
        WebRTC 风格的噪声估计模块
        
        参数:
        fft_size: 频域频谱的大小 (n_bins = fft_size//2 + 1)
        sr: 采样率 (默认为16000Hz)
        bands_config: 频带配置 ("webRTC", "mel", 或 "linear")
        """
        # 频谱元数据
        self.n_bins = fft_size // 2 + 1
        self.sr = sr
        
        # 噪声估计参数
        self.ALPHA = 0.9  # 先验信噪比更新的平滑因子
        self.MIN_ENERGY = 1e-10  # 避免除以零的最小能量值
        
        # 频带配置
        self._setup_frequency_bands(bands_config)
        
        # 初始化噪声估计状态
        self.reset()
    
    def reset(self):
        """重置噪声估计器的内部状态"""
        # 噪声功率谱估计
        self.noise_psd = np.zeros(self.n_bins) + self.MIN_ENERGY
        
        # 帧级状态
        self.frame_count = 0
        self.last_posterior_snr = np.zeros(self.n_bins)
        self.last_prior_snr = np.zeros(self.n_bins)
        
        # 频带级状态
        self.min_psd = np.zeros(self.n_bands) + np.inf  # 频带最小功率
        self.min_persistence = np.zeros(self.n_bands)  # 最小值的持续时间计数器
        self.psd_estimate = np.zeros(self.n_bands)  # 频带噪声估计
        self.psd_smooth = np.zeros(self.n_bands)  # 平滑的功率估计
    
    def _setup_frequency_bands(self, config: str):
        """设置频率带分组方案"""
        # 频率轴 (0Hz 到 Nyquist)
        freqs = np.linspace(0, self.sr/2, self.n_bins)
        
        if config == "webRTC":
            # WebRTC的分频方案 (调整边界不超过Nyquist频率)
            possible_edges = [0, 100, 200, 300, 400, 500, 600, 700, 800, 
                            900, 1000, 2000, 3000, 4000, 8000]
            self.band_edges = [f for f in possible_edges if f <= freqs[-1]]
        elif config == "mel":
            # Mel频率标度分组 (不超过n_bins)
            n_bands = min(20, self.n_bins//2)
            self.band_edges = np.linspace(0, self.sr/2, n_bands+1)
        else:  # linear
            # 线性分组 (不超过n_bins)
            n_bands = min(10, self.n_bins//2)
            self.band_edges = np.linspace(0, self.sr/2, n_bands+1)
        
        # 确保最后一个边界不超过Nyquist频率
        self.band_edges[-1] = min(self.band_edges[-1], freqs[-1])
        
        # 创建频带索引映射
        self.band_indices = []
        for i in range(len(self.band_edges)-1):
            start_freq = self.band_edges[i]
            end_freq = self.band_edges[i+1]
            indices = np.where((freqs >= start_freq) & (freqs < end_freq))[0]
            if len(indices) == 0:
                # 如果频带为空，使用前一个频带的最后一个bin
                if i > 0 and len(self.band_indices[-1]) > 0:
                    indices = np.array([self.band_indices[-1][-1]])
                else:
                    indices = np.array([0])  # 最低频bin
            self.band_indices.append(indices)
        
        self.n_bands = len(self.band_indices)
    
    def update_band_psd(self, band_psd: List[float]):
        """
        更新每个频带的噪声估计 (核心算法)
        
        参数:
        band_psd: 每个频带的功率谱密度估计
        """
        # 更新每个频带的最小功率估计
        for b in range(self.n_bands):
            # 当前频带的平滑功率估计
            beta = 0.7  # 平滑因子
            self.psd_smooth[b] = beta * self.psd_smooth[b] + (1 - beta) * band_psd[b]
            
            # 寻找当前频带的最小功率估计
            if self.psd_smooth[b] < self.min_psd[b]:
                self.min_psd[b] = self.psd_smooth[b]
                self.min_persistence[b] = 0
            else:
                self.min_persistence[b] += 1
                
            # 如果最小值太旧，则重置它
            if self.min_persistence[b] > 100:  # 约1秒 (100帧)
                self.min_psd[b] = self.psd_smooth[b]
                self.min_persistence[b] = 0
            
            # 更新该频带的噪声估计
            alpha = 0.05
            self.psd_estimate[b] = alpha * self.min_psd[b] + (1 - alpha) * self.psd_estimate[b]
    
    def estimate_noise(self, power_spectrum: np.ndarray) -> np.ndarray:
        """
        估计当前帧的噪声功率谱 (主要入口函数)
        
        参数:
        power_spectrum: 当前帧的功率谱 (shape: [n_bins])
        
        返回:
        当前帧更新后的噪声功率谱估计
        """
        # 确保功率谱长度匹配
        assert len(power_spectrum) == self.n_bins, \
            f"功率谱长度{len(power_spectrum)} != n_bins{self.n_bins}"
        
        self.frame_count += 1
        
        # 1. 为每个频带计算平均功率
        band_psd = []
        for indices in self.band_indices:
            # 确保索引在有效范围内
            valid_indices = indices[indices < len(power_spectrum)]
            if len(valid_indices) == 0:
                band_psd.append(0.0)
            else:
                band_psd.append(np.mean(power_spectrum[valid_indices]))
        
        # 2. 更新频带级噪声估计
        self.update_band_psd(band_psd)
        
        # 3. 构建全频带噪声估计
        noise_psd = np.zeros_like(power_spectrum)
        for b, indices in enumerate(self.band_indices):
            for idx in indices:
                noise_psd[idx] = self.psd_estimate[b]
        
        # 4. 递归平滑更新噪声估计
        alpha = 0.7  # 噪声估计的平滑因子
        self.noise_psd = alpha * self.noise_psd + (1 - alpha) * noise_psd
        
        # 5. 确保最小值保护
        self.noise_psd = np.maximum(self.noise_psd, self.MIN_ENERGY)
        
        # 6. 计算并保存SNR用于后续处理
        self.last_posterior_snr = power_spectrum / self.noise_psd
        self.last_prior_snr = self._update_prior_snr(self.last_posterior_snr)
        
        return self.noise_psd.copy()
    
    def _update_prior_snr(self, posterior_snr: np.ndarray) -> np.ndarray:
        """
        更新先验信噪比 (SNR) 使用决策导向方法
        
        参数:
        posterior_snr: 当前帧的后验信噪比
        
        返回:
        更新后的先验信噪比
        """
        # 对于第一帧，使用后验SNR作为先验SNR
        if self.frame_count == 1:
            return posterior_snr
        
        # 决策导向递归更新
        # gain = 先验SNR/(1+先验SNR) 是前一帧的Wiener滤波器增益
        prev_gain = self.last_prior_snr / (1 + self.last_prior_snr)
        prior_snr = self.ALPHA * (prev_gain ** 2) * self.last_posterior_snr + (1 - self.ALPHA) * np.maximum(posterior_snr - 1, 0)
        
        return prior_snr
    
    def get_noise_estimate(self) -> np.ndarray:
        """获取当前噪声功率谱估计"""
        return self.noise_psd.copy()
    
    def get_prior_snr(self) -> np.ndarray:
        """获取当前帧的先验信噪比估计"""
        return self.last_prior_snr.copy()
    
    def get_posterior_snr(self) -> np.ndarray:
        """获取当前帧的后验信噪比估计"""
        return self.last_posterior_snr.copy()


# 测试代码
if __name__ == "__main__":
    # 假设我们有256点FFT (n_bins = 129)
    fft_size = 256
    noise_estimator = NoiseEstimator(fft_size, bands_config="webRTC")
    
    # 模拟一些帧的处理
    print("模拟噪声估计处理...")
    for i in range(10):
        # 生成模拟功率谱 (包含噪声)
        n_bins = fft_size // 2 + 1
        true_noise = np.random.rand(n_bins) * 0.1 + 0.01
        signal_power = np.random.rand(n_bins) * 0.5 + 0.1 + true_noise
        
        # 估计噪声
        noise_est = noise_estimator.estimate_noise(signal_power)
        
        print(f"帧 {i+1}: 最大噪声估计 = {np.max(noise_est):.4f}, 平均噪声估计 = {np.mean(noise_est):.4f}")
    
    # 获取最终噪声估计和SNR指标
    final_noise = noise_estimator.get_noise_estimate()
    prior_snr = noise_estimator.get_prior_snr()
    posterior_snr = noise_estimator.get_posterior_snr()
    
    print("\n噪声估计统计:")
    print(f"最小值: {np.min(final_noise):.6f}, 最大值: {np.max(final_noise):.4f}, 均值: {np.mean(final_noise):.4f}")
    
    print("\nSNR统计:")
    print(f"先验SNR范围: {np.min(prior_snr):.2f} - {np.max(prior_snr):.2f}")
    print(f"后验SNR范围: {np.min(posterior_snr):.2f} - {np.max(posterior_snr):.2f}")